'''
This module contains the peer protocol used to share or download files
in the simple-torrent system.

@author: Hatem Oraby
'''


from os import path,rename
from warnings import warn
from hashlib import sha1,md5
from random import randint

from twisted.internet.protocol import ServerFactory, Protocol, ClientCreator
from twisted.internet.defer import Deferred
from twisted.internet import reactor

import torrentFile
from peerToTracker import registerWithTracker


announceTrackerTimeout = 4.75*60#Announce self to tracker


class serverStates:
    NothingYet = 0
        
class serverCmds:
    getPiece = "GETPIECE"


class serverRspns:
    pieceExist = "PIECEEXIST"
    pieceNotExist = "PIECENOTEXIST"


class TorrentManager:
    def __init__(self,torrentName, torrentPath):
        self._torrent = torrentFile.parseTorrent(torrentName, torrentPath)
        self._peers = set()
        self._piecesOwn = []
        self._piecesCount = self._torrent.getPiecesCount()
        self._chunkSizeKB = self._torrent.getChunkSize()*1024
        self._fileSize = self._torrent.getFileSize()
        self._fileName = self._torrent.getFileName()
        #Know if I'm a seeder or a leecher.
        if path.isfile(self._fileName):
            self._isLeecher = False
            #Then I have all the pieces
            for i in range(self._piecesCount):
                self._piecesOwn.append(True)
            #In the future I may do extra checking on the file hashes 
            #to ensure the integrity of the file
            #self._fileHandler = open(fileName, 'rb')
            self._piecesComplete = self._piecesCount
        else:
            self._isLeecher = True
            for i in range(self._piecesCount):
                self._piecesOwn.append(False)
            self._tempName = "_" + self._fileName
            fileHandler = open(self._tempName, 'w+b') #r+b: Read and write binary
            #Allocate space in the file that we just created
            remainingBytes = self._fileSize
            while remainingBytes:
                fetchSize = min(self._chunkSizeKB, remainingBytes)
                data = "0"*fetchSize
                fileHandler.write(data)
                remainingBytes -= fetchSize
            fileHandler.close()
            self._piecesComplete = 0
        ip, port = self._torrent.getTrackerUrl().rsplit(":",1)
        self._trackerAddr = ip , int(port.strip())
        self._piecesHash = self._torrent.getPiecesHash()
    
    
    def addPeer(self,peer):
        self._peers.add(peer)
    
    def removePeer(self, peer):
        self._peers.remove(peer)
    
    def getPeers(self):
        return list(self._peers)
    
    def isPieceExist(self, index):
        return index in range(self._piecesCount)
    
    def isPieceDownloaded(self, index):
        return self._piecesOwn[index]
    
    def putPiece(self, index, chunk):
        #Do checkings here
        if not self.isLeecher():
            raise "I'am a seeder, not leecher"
        if not self.isPieceExist(index):
            raise "Piece Index out of bound"
        
        chunkHash = sha1(chunk).hexdigest()
        if chunkHash != self._piecesHash[index]:
            print "Incoming Hash: " , chunkHash
            print "Expected Hash: " , self._piecesHash[index]
#            for xHash in self._piecesHash:
#                print "-Expected Hash: " , xHash
            raise "Piece Corrupted"#Should return with appropriate error
        if self.isPieceDownloaded(index):
            #warn("Piece already exist. Skipping it.")
            return 
        handler = self._getFileHandler()
        offset = self._chunkSizeKB*index
        handler.seek(offset)
        handler.write(chunk)
        self._piecesOwn[index] = True
        self._piecesComplete += 1
        if self._piecesComplete == self._piecesCount:
            #print "File Transfer Complete. Checking file hash"
            MD5 = md5()
            chunk = max(self._chunkSizeKB, 5*1024*1024)#read in chunk of 5MBs
            handler.seek(0)
            while True:
                data = handler.read(chunk)
                if not data:
                    break
                MD5.update(data)
            #for i in self._piecesOwn:
            #    print i
            if MD5.hexdigest() != self._torrent.getFileHash():
                raise "Corrupted File Hash. Is it corrupted torrent file??."
            else:
                #print "MD5 File Check Passed."
                #self.printBitmap()
                fileName = self._torrent.getFileName()
                tempFileName = handler.name
                handler.close()
                rename(tempFileName, fileName)
                self._isLeecher = False
        handler.close()
        
    def getPiece(self, index):
        handler = self._getFileHandler()
        if not self.isPieceExist(index):
            raise "Piece out of range."
        offset = index*self._chunkSizeKB#*1024
        handler.seek(offset)
        data = handler.read(self._chunkSizeKB)
        handler.close()
        return data
    
    def _getFileHandler(self):
        #Create a new handler for each new operation,
        #this way we are sure that we are thread safe.
        if self._isLeecher:
            return  open(self._tempName, 'r+b') #r+b: Read and write binary
        else:
            return  open(self._fileName, 'rb')
        
    def isLeecher(self):
        return  self._isLeecher
    
    def getTrackerAddr(self):
        return self._trackerAddr
    
    def getTotalNumPieces(self):
        return self._piecesCount
    
    def printBitmap(self):        
        for i in self._piecesOwn:
            if i:
                print "#",
            else:
                print "_",
        print ""

class requestssHandler(Protocol):
    def __init__(self, torrentManager):      
        self._torrentManager = torrentManager

    def pieceRequest(self,data):
        try:
            command, params = data.split(":",1)
        except ValueError:
            #Wrong command sent
            print   "E51: Server received wrong format"
            self.transport.loseConnection()
        if command.strip().upper() != serverCmds.getPiece:
            print "E52: Unexpected command."
            print "Expected: " ,serverCmds.getPiece, "\tGot: ",command
            self.transport.loseConnection()
        #Command is correct
        params = params.strip()
        try:
            pieceIndex = int(params)
        except ValueError:
            print "E52: Wrong param. Expected int got: ",params
            self.transport.loseConnection()
        #Check first that the piece exist
        if not self._torrentManager.isPieceExist(pieceIndex):
            print "E53: Wrong Piece index."
            print "\tGot: ",pieceIndex
            self.transport.loseConnection()
        #Param and PieceIndex are correct
        #If i have the chunk I should answer with OK and send the
        #chunk, else I should point out that I don't have it.
        if not self._torrentManager.isPieceDownloaded(pieceIndex):
            response = serverRspns.pieceNotExist + ":\n"
            self.transport.write(response)
            self.transport.loseConnection()
        else:
            chunk = self._torrentManager.getPiece(pieceIndex)
            #print "Chunk length: " , len(chunk)
            #print "Piece Index: ", pieceIndex
            #print "Chunk Hash: " , sha1(chunk).hexdigest()
            response = serverRspns.pieceExist + ":" + str(len(chunk)) + "\n"
            response += chunk
            self.transport.write(response)
            #Let twisted handle splitting the data if its big.
            self.transport.loseConnection()
        
    def dataReceived(self,data):
        if data.startswith(serverCmds.getPiece):
            self.pieceRequest(data)
        else:
            print "50:Bad Command sent. Got:" + data
            self.transport.loseConnection()


class requestsHandlerFactory(ServerFactory):
    protocol = requestssHandler
    
    def __init__(self,torrentManager):
        self._torrentManager = torrentManager
        
    def buildProtocol(self, addr):
        p = self.protocol(self._torrentManager)
        p.factory = self
        return p


class PieceLeecher(Protocol):
    class state:
        nothingYet = 0
        recievingData = 1
        done = 2
        
    def __init__(self, torrentManager, defer):
        self._torrentManager = torrentManager
        self._state = PieceLeecher.state.nothingYet
        self._dataReceived = ""
        self._defer = defer
        self._dataSize = 1*1024*1024*1024
        #print "Initializing pieceLeacher"

    def _handleInitial(self,data):
        try:
            command, params_n_data = data.split(":",1)
        except ValueError:
            #Wrong command sent
            print   "E61: Server received wrong format"
            self._defer.errback("Server bad format")
            self.transport.loseConnection()
        
        if command != serverRspns.pieceExist:
            raise  "E62: Something bad happened on my side."
        
        params, data = params_n_data.split("\n",1)
        try:
            self._dataSize = int(params.strip())
        except ValueError:
            print "E63: Server sent bad param. Expected int Got: " + params
            self._defer.errback("Bad Params")
        self._state = PieceLeecher.state.recievingData
        self._pushData(data)
        
    def _pushData(self,data):
        self._dataReceived += data
        if len(self._dataReceived) > self._dataSize:
            warn("E64: I Over-received data???")
            self.transport.loseConnection()
        elif len(self._dataReceived) == self._dataSize:
            #Should i Wait for him to disconnect or should I disconnect??
            self.transport.loseConnection()
            #Connection lost will call handle_finishing()

    def dataReceived(self,data):
        if self._state == PieceLeecher.state.nothingYet:
            if data.startswith(serverRspns.pieceExist):
                self._handleInitial(data)
            elif data.startswith(serverRspns.pieceNotExist):
                #Is it really an error??
                self._defer.errback(Exception("Piece Doesn't exist"))
                self._state = PieceLeecher.state.done
                self.transport.loseConnection()
            else:
                print "60:Bad Command Received. Got:" + data
                self._defer.errback(Exception("Back server Command"))
                self.transport.loseConnection()
        else:
            self._pushData(data)

    def connectionMade(self):
        #print "Piece Leecher connection made."
        if not self._torrentManager.isLeecher():
            self.transport.loseConnection()
        numofPieces = self._torrentManager.getTotalNumPieces()
        index = randint(0,numofPieces -1)
        #If threads were used, can the next part block forever??
        #well it's better to re-check that I'm still a leecher.
        while self._torrentManager.isPieceDownloaded(index) \
            and self._torrentManager.isLeecher():
            index = randint(0, numofPieces -1)
        msg = serverCmds.getPiece + ":" + str(index)
        self._indexRequested = index
        self.transport.write(msg)

    def connectionLost(self, reason):
        self._handleFinishing()
        
    def _handleFinishing(self):
        if len(self._dataReceived) > self._dataSize:
            print "E65: Can I be Over-recieve data or be called twice???"
            self._defer.errback("over-received data")
        elif len(self._dataReceived) == self._dataSize:
            #Should i Wait for him to disconnect or should I disconnect??
            self._torrentManager.putPiece(self._indexRequested, self._dataReceived)
            #I should handle torrentManager responses.
            self._defer.callback(self._indexRequested)
        else:
            if self._state != PieceLeecher.state.done:
                msg = "Connection Terminated or lost before transfer Complete."
                print "Data size: ", len(self._dataReceived)
                print "Expected: " , self._dataSize
                self._defer.errback(Exception(msg))
            
    
class Leecher():
    
    def __init__(self,torrentManager):
        self._torrentManager = torrentManager
        self._fetchTimeout = 60#1 min
        self._fetchInProgress = 0
        self._fetchLimit = 1
        self._tryGetPeersCount = 0#Those two variables are just
        self._tryGetPeersLimit = 20# for printing warning messages 
        self._tryLeech()
        
    def _tryLeech(self):
        if self._torrentManager.isLeecher():
            peers = self._torrentManager.getPeers()
            if not len(peers):
                self._tryGetPeersCount += 1
                if self._tryGetPeersCount >= self._tryGetPeersLimit:
                    #warn("Got No Peers, Waiting for new peers in next "\
                    #     "tracker registration round")
                    self._tryGetPeersCount = 0
                reactor.callLater(self._fetchTimeout,self._tryLeech)
                #self._getPeers()
            else:
                #print "Leeching again."
                #Pick up a peer randomly, be thread safe through isLeeacher()
                while self._fetchInProgress < self._fetchLimit \
                        and self._torrentManager.isLeecher():
                    peerIndex = randint(0,len(peers) -1)
                    #print "Peer Index: " , peerIndex
                    peer = peers[peerIndex]
                    defer = Deferred()
                    defer.addCallbacks(self._peerDone,self._peerDoneErr
                                       ,errbackArgs = [peer])
                    c = ClientCreator(reactor, PieceLeecher, 
                                      self._torrentManager,defer)
                    dConn = c.connectTCP(peer.getIP(), peer.getPort())
                    dConn.addErrback(self._peerConnFailed, peer)
                    self._fetchInProgress += 1
        
    def _peerDone(self, something):
        #print "Connection done with peer."
        self._fetchInProgress -= 1
        self._tryLeech()
    
    def _peerDoneErr(self, reason, peer):
        #print "Failed to communicate with peer. Reason: " 
        #print  reason
        self._torrentManager.removePeer(peer)
        self._peerDone(reason)
        
    def _peerConnFailed(self,reason, peer):
        #print "Failed to connect to peer: ",
        #print peer.getIP(), ":", peer.getPort()
        try:
            self._torrentManager.removePeer(peer)
        except ValueError:
            #Maybe I did previously remove him.
            pass
        self._fetchInProgress -= 1
        self._tryLeech()

def registerSelfCont(torrentManager, port, registerTimeout):
    
    def registeredSucceffuly(peers):
        #print peers
        for peer in peers:
            torrentManager.addPeer(peer)
        reactor.callLater(registerTimeout, registerSelfCont, 
                          torrentManager, port, registerTimeout)
        
    def failedToRegister(reason):
        msg = "Something bad happened while registering with tracker.\n" 
        msg += "Stopped Trying to register"
        warn(msg)
    
    def connectionTrackerFailed(reason):
        warn("Failed To connect to tracker. Failed to register.")
    
    trackerAddr = torrentManager.getTrackerAddr()
    dRegister = Deferred()
    dRegister.addCallbacks(registeredSucceffuly, failedToRegister)    
    d = registerWithTracker(trackerAddr, port, dRegister)
    d.addErrback(connectionTrackerFailed)

def recursivePrintMap(torrentManager):
    if torrentManager.isLeecher():
        torrentManager.printBitmap()
        reactor.callLater(10,recursivePrintMap,torrentManager)

def main():
    #Parse sys.argv
    from sys import argv
    if len(argv) != 2:
        print "Wrong usage."
        print "Usage: TORRENTFILE.tor TrackerIP"
        exit(-1)
    torrneFileName = argv[1]
    
    torrentManager = TorrentManager(torrneFileName,"")
    factory = requestsHandlerFactory(torrentManager)
    twistedPort = reactor.listenTCP(0, factory)
    port = twistedPort._realPortNumber#Get the port that we are listening on
    
    registerSelfCont(torrentManager, port, announceTrackerTimeout)
    if torrentManager.isLeecher():
        #print "start doing the leeching stuff"
        Leecher(torrentManager)
    #recursivePrintMap(torrentManager)
    reactor.run()
    
if __name__ == '__main__':
    main()